'''
将3D的nii数据，切成2D切片。每张切片的名称包含其来源于哪张nii，位于nii的第几张切片
'''
import os
import sys

from scipy import ndimage

sys.path.append(os.path.split(sys.path[0])[0])
import SimpleITK as sitk
from multiprocessing.dummy import Pool
from config import zunYiParameter as para
import numpy as np

def mynormalize(img:np.ndarray):
    return (img-img.mean())/img.std()

def identicShape(volume_array,new_shape,volume_or_seg):

    old_shape=volume_array.shape
    dimension_adjustment=new_shape/old_shape
    if volume_or_seg=="volume":
        volume_array_new = ndimage.zoom(volume_array, dimension_adjustment, order=3)
    elif volume_or_seg=="seg":
        volume_array_new = ndimage.zoom(volume_array, dimension_adjustment, order = 0)

    return volume_array_new

def process(file):
    # 将CT和标签入读内存
    print(file)
    volume = sitk.ReadImage(os.path.join(para.nii_volume_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)

    volume_array=mynormalize(volume_array)#以每个3D为对象进行归一化
    #print('处理前的volume shape:',volume_array.shape,end = '   ')
    seg = sitk.ReadImage(os.path.join(para.nii_seg_path, file.replace('volume', 'segmentation')), sitk.sitkUInt8)
    seg_array = sitk.GetArrayFromImage(seg)

    #统一维度
    volume_array=identicShape(volume_array, new_shape = np.array([volume_array.shape[0],384,384])
                              ,volume_or_seg = "volume")
    seg_array = identicShape(volume_array, new_shape = np.array([volume_array.shape[0], 384, 384])
                                , volume_or_seg = "seg")
    print(volume_array.shape,seg_array.shape)
    seg_file_name=file.replace('volume', 'segmentation').split('.')[0]
    ct_file_name=file.split('.')[0]

    # seg_array=processLiverOrTumorMask(para.Liver_or_Tumor,seg_array)
    # 将灰度值在阈值之外的截断掉
    # volume_array[volume_array > para.upper] = para.upper
    # volume_array[volume_array < para.lower] = para.lower

    all_indexs = [i for i in range(len(volume_array))]
    #保存图片
    file_names=[]

    for index in all_indexs:
        ct_name=os.path.join("volume",ct_file_name+"_"+str(index)+'.nii.gz')
        seg_name=os.path.join("segmentation", seg_file_name+ "_" + str(index)+'.nii.gz' )
        file_names.append(ct_name+" "+seg_name)
        #保存
        ct_path = os.path.join(para.cut2d_save_path,ct_name)
        seg_path = os.path.join(para.cut2d_save_path,seg_name)
        sitk.WriteImage(sitk.GetImageFromArray(volume_array[index]),ct_path)
        sitk.WriteImage(sitk.GetImageFromArray(seg_array[index]), seg_path)
    return file_names


def main():
    if not os.path.exists(para.cut2d_save_path_vol):
        os.makedirs(para.cut2d_save_path_vol)
    if not os.path.exists(para.cut2d_save_path_seg):
        os.makedirs(para.cut2d_save_path_seg)

    if not os.path.exists(os.path.dirname(para.test2d_id_path)):
        os.makedirs(os.path.dirname(para.test2d_id_path))


    volume_name_list = os.listdir(para.nii_volume_path)
    print("process nii file number:", len(volume_name_list))

    #处理测试集
    pool_for_test = Pool(2)
    result=pool_for_test.map(process, volume_name_list)
    file_names_flatten = []
    for i in range(len(result)):
        file_names_flatten += result[i]
    print("切片共{}张".format(str(len(file_names_flatten))))
    # test_file_names_flatten = np.array(test_file_names_flatten)
    # print("test size", test_file_names_flatten.shape)
    # pool_for_test.close()
    # pool_for_test.join()
    #
    #
    # np.savetxt(para.test2d_id_path, test_file_names_flatten, fmt = "%s")

if __name__ == '__main__':
    '''
    test size (768,)
    '''
    main()